package priv.tuyou.redis.limit;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;

/**
 * 令牌桶-限流器
 *
 * 按一定速率生成令牌到一个池子中，也就是增加incr key
 * @author tuyou_cd
 * @date 2021/4/23 15:04
 */
public class TokenBucketRateLimiter extends AbstractRateLimiter implements Runnable {

    /**
     * 是否初始化完成
     */
    volatile boolean initialed;

    /**
     * 生成令牌的速度
     */
    int speed;

    public TokenBucketRateLimiter(JedisPool pool, int period, int limit) {
        super(pool, period, limit);
        speed = period * 1000 / limit;

        //启动周期性生成令牌
        new Thread(this).start();
    }

    @Override
    public boolean visit(int count) {
        if (!initialed) {
            return false;
        }

        try (Jedis jedis = pool.getResource()) {
            String s = jedis.get(key);
            if (s == null) {
                return false;
            }
            if (Integer.parseInt(s) > 0) {
                Long decr = jedis.decr(key);
                if (decr < 0) {
                    return false;
                }
            }
        }
        return true;
    }

    @Override
    public void run() {
        //初始化令牌桶
        try (Jedis jedis = pool.getResource()) {
            jedis.setnx(key, "0");
        }
        initialed = true;

        while (true) {
            try (Jedis jedis = pool.getResource()) {
                jedis.incr(key);
            }
            try {
                Thread.sleep(speed);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}
